#!/usr/bin/env python
# coding: utf-8

"""
export_rddhw.py
create by cuiyu 2025-03-06
说明：
 * 导出 onnx 和 tensorrt 模型文件
 * 参数：见 modelExport() 函数中的参数设置
 * 导出模型大约需要 7分30秒；绝大部分时间用在了 tensorrt 模型转换上
 * 导出的模型文件保存在输入模型文件的相同目录下

用法:
  export_rddhw.py

更新:
2025-03-06:
"""

import sys
from pathlib import Path
from ultralytics import YOLO

APP_FILE = Path(sys.argv[0])
APP_ROOT = APP_FILE.parent

def modelExport(model_file: Path, imgsz):
    """
    导出模型文件

    参数：
        model_file: 模型文件路径
        imgsz:      输入图像尺寸，可以是整数或 (height, width) 元组，最好和训练尺寸一致
    """
    model = YOLO(model_file)
    out_file = model.export(
        format = "engine",      # engine=TensorRt, onnx=ONNX
        imgsz = imgsz,          # 2048 输入图像尺寸（最好和训练尺寸一致）
        batch = 1,              # 固定batch size为1以获得最佳性能
        device = 0,             # 使用 GPU 0 进行导出，要使用CPU则设置为 "cpu"
        half = True,            # 启用 FP16 量化加速推理 启用半精度(FP16)量化
        workspace = 6,          # 设置TensorRT工作空间大小为 xGB ，可根据硬件调整
        simplify = True,        # 使用 onnxslim 简化 ONNX 导出的模型图，可能提高性能和兼容性。（TensorRT 需要先生成 ONNX）
        dynamic = False,        # 允许动态输入尺寸。固定尺寸更快，但是需要在推理时调整输入尺寸。dynamic=True可能导致TensorRT转换失败
        int8 = False,           # 可选INT8量化（需要校准数据）
    )
    print(f"model export to {out_file}")

    # 导出 tensorrt 模型在 onnx -> tensorrt 需要的时间比较长，也可能会失败，失败后可还有手动转换。
    # 手动转换的命令：
    # trtexec --onnx=/aaby/netdisk/mydata/ai-data/model/ultralytics/seg/yolo11n-seg.onnx --saveEngine=/aaby/netdisk/mydata/ai-data/model/ultralytics/seg/yolo11n-seg.engine --explicitBatch

    pass

def _main():
    from runtimer import runtimer_create

    #==== 参数设置 =======================
    model_file = Path("/aaby/netdisk/mydata/ai-data/model/jgjc/jgjc-20250805-181212/weights/best.pt")
    imgsz = (640, 640)
    #====================================

    # 模型转换耗时较长，大约需要 7分钟
    with runtimer_create("模型转换"):
        modelExport(model_file, imgsz)

if __name__ == "__main__":
    # logging.basicConfig(level=logging.DEBUG) 
    _main()
